package net.oschina.esb.util;

import org.apache.commons.lang3.StringUtils;
import sun.net.spi.nameservice.NameService;
import sun.net.util.IPAddressUtil;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.lang.reflect.Field;
import java.net.InetAddress;
import java.net.UnknownHostException;
import java.util.*;

/**
 * @author Arvin
 * @time 2018/5/4 22:31
 */
public class DnsUtil {

    /** 是否启用自定义解析 DNS */
    private static volatile boolean enabledCustomDns = false;

    /** 域名 --> ip 列表集合 */
    private static Map<String, SortedSet<String>> DOMAIN_IP_MAP = new HashMap<>();

    /** 自定义的 HOST */
    private static final NameService CUSTOM_DNS_NAME_SERVICE = new NameService() {
        @Override
        public InetAddress[] lookupAllHostAddr(String host) throws UnknownHostException {

            if (DOMAIN_IP_MAP == null || DOMAIN_IP_MAP.isEmpty()) {
                throw new UnknownHostException(host);
            }

            SortedSet<String> ips = DOMAIN_IP_MAP.get(host);
            if (null == ips || ips.isEmpty()) {
                throw new UnknownHostException(host);
            }

            InetAddress[] arrays = new InetAddress[ips.size()];
            int index = 0;
            for (String ip : ips) {
                arrays[index++] = InetAddress.getByAddress(host, InetAddress.getByName(ip).getAddress());
            }

            return arrays;
        }

        @Override
        public String getHostByAddr(byte[] bytes) throws UnknownHostException {
            throw new UnknownHostException(new String(bytes));
        }
    };

    /** 最后一次被覆盖的自定义缓存 */
    private static Set<String> LAST_OVERRIDE_CUSTOM_DOMAINS = new HashSet<>();

    /**
     * 清理 自定义的 DNS 缓存
     */
    public static void clearCustomDnsCache() {
        if (LAST_OVERRIDE_CUSTOM_DOMAINS == null) {
            LAST_OVERRIDE_CUSTOM_DOMAINS = new HashSet<>();
        }
        if (!DOMAIN_IP_MAP.isEmpty()) {
            LAST_OVERRIDE_CUSTOM_DOMAINS.addAll(DOMAIN_IP_MAP.keySet());
        }
        if (LAST_OVERRIDE_CUSTOM_DOMAINS.isEmpty()) {
            return;
        }

        Map<String, Object> cacheMap = getAddressCache();
        if (cacheMap == null || cacheMap.isEmpty()) {
            return;
        }

        for (String domain : LAST_OVERRIDE_CUSTOM_DOMAINS) {
            cacheMap.remove(domain);
        }
    }


    /**
     * 清理 域名的 DNS 缓存
     */
    public static void clearDnsCache(String domain) {

        if (StringUtils.isBlank(domain)) {
            return;
        }

        Map<String, Object> cacheMap = getAddressCache();
        if (cacheMap == null || cacheMap.isEmpty()) {
            return;
        }
        cacheMap.remove(domain);
    }

    /**
     * 清理 域名的 DNS 缓存
     *
     * @param domains 要清理缓存的 域名列表
     */
    public static void clearDnsCache(Collection<String> domains) {

        if (null == domains || domains.isEmpty()) {
            return;
        }

        Map<String, Object> cacheMap = getAddressCache();
        if (cacheMap == null || cacheMap.isEmpty()) {
            return;
        }
        for (String domain : domains) {
            cacheMap.remove(domain);
        }
    }

    /**
     * 添加自定义的域名解析
     *
     * @param domain 域名
     * @param ip     IP
     */
    public synchronized static void addOrUpdateCustomDns(String domain, String ip) {

        if (StringUtils.isAnyBlank(domain, ip) || !isValidIp(ip)) {
            throw new RuntimeException("域名IP不合法： " + ip + " " + domain);
        }

        if (isCustomDnsEnabled()) {
            SortedSet<String> ips = new TreeSet<>();
            ips.add(ip);
            DOMAIN_IP_MAP.put(domain, ips);
            clearDnsCache(domain);
        } else {
            loadCustomDnsByLines(new ArrayList<>(Arrays.asList(ip + " " + domain)));
        }

    }

    /**
     * 添加或者更新自定义DNS解析
     *
     * @param lines ip 域名行， 如 IP DOMAIN
     * @return 返回域名解析MAP
     */
    public synchronized static Map<String, List<String>> addOrUpdateCustomDns(List<String> lines) {

        Map<String, SortedSet<String>> domainIpMap = buildCustomDnsMapByLines(lines);

        if (null != domainIpMap && !domainIpMap.isEmpty()) {
            DOMAIN_IP_MAP.putAll(domainIpMap);
            enabledCustomDns();
        }

        return copyDomainIpMap(domainIpMap);
    }

    /**
     * 删除自定义的 DNS 缓存
     *
     * @param domain 域名
     */
    public synchronized static void removeCustomDns(String domain) {
        if (StringUtils.isNotBlank(domain) && DOMAIN_IP_MAP != null && DOMAIN_IP_MAP.containsKey(domain)) {
            DOMAIN_IP_MAP.remove(domain);
            clearDnsCache(domain);
        }
    }

    /**
     * 删除自定义的 DNS 缓存
     *
     * @param domains 域名
     */
    public synchronized static void removeCustomDns(Collection<String> domains) {

        if (null != domains && !domains.isEmpty()) {
            for (String domain : domains) {
                DOMAIN_IP_MAP.remove(domain);
            }
            clearDnsCache(domains);
        }

    }

    /**
     * 获取 InetAddress DNS 缓存
     *
     * @return 返回 缓存
     */
    protected static Map<String, Object> getAddressCache() {
        try {
            final Field cacheField = InetAddress.class.getDeclaredField("addressCache");
            cacheField.setAccessible(true);
            final Object addressCache = cacheField.get(InetAddress.class);

            Class clazz = addressCache.getClass();
            final Field cacheMapField = clazz.getDeclaredField("cache");
            cacheMapField.setAccessible(true);
            return (Map) cacheMapField.get(addressCache);
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage(), e);
        }
    }

    /**
     * 自定义 DNS nameService
     */
    private static void customDnsNameService() {
        if (isCustomDnsEnabled()) {
            Class<?> clazz = InetAddress.class;

            Field field = null;
            try {
                field = clazz.getDeclaredField("nameServices");

                field.setAccessible(true);

                Object obj = field.get(clazz);

                if (obj instanceof List) {
                    List<NameService> nameServiceList = (List<NameService>) obj;

                    boolean hadConfig = false;
                    for (int i = 0; i < nameServiceList.size(); ++i) {
                        Object nameService = nameServiceList.get(i);
                        if (nameService == CUSTOM_DNS_NAME_SERVICE) {
                            hadConfig = true;
                            break;
                        }
                    }

                    if (!hadConfig) {
                        nameServiceList.add(0, CUSTOM_DNS_NAME_SERVICE);
                    }

                }

                // 清理自定义的本地缓存
                clearCustomDnsCache();

            } catch (Exception e) {
                e.printStackTrace();
            } finally {
                if (field != null) {
                    field.setAccessible(false);
                }
            }
        }
    }

    /**
     * 禁用 自定义的 DNS 缓存
     */
    public static void disabledCustomDns() {
        enabledCustomDns = false;
        // 删除 NameService
        Class<?> clazz = InetAddress.class;

        Field field = null;
        try {
            field = clazz.getDeclaredField("nameServices");

            field.setAccessible(true);

            Object obj = field.get(clazz);

            boolean hadRemove = false;
            if (obj instanceof List) {
                List<NameService> nameServiceList = (List<NameService>) obj;

                for (int i = 0; i < nameServiceList.size(); ++i) {
                    Object nameService = nameServiceList.get(i);
                    if (nameService == CUSTOM_DNS_NAME_SERVICE) {
                        nameServiceList.remove(i);
                        hadRemove = true;
                    }
                }
            }

            if (hadRemove) {
                LAST_OVERRIDE_CUSTOM_DOMAINS = new HashSet<>();
                clearCustomDnsCache();
            }

        } catch (Exception e) {
            e.printStackTrace();
        } finally {
            if (field != null) {
                field.setAccessible(false);
            }
        }
    }

    public static void enabledCustomDns() {
        enabledCustomDns = true;
        customDnsNameService();
    }

    public static boolean isCustomDnsEnabled() {
        return enabledCustomDns;
    }

    /**
     * 加载自定义的 DNS，要求输入的文件格式是一行表是一个 IP 域名， 可以使用 # 注释
     * <p>
     * IP DOMAIN
     * # 注释
     *
     * @param inputStream 文件输入流
     */
    public synchronized static Map<String, List<String>> loadCustomDns(InputStream inputStream) throws Exception {

        if (null == inputStream) {
            throw new RuntimeException("自定义DNS输入流为空，无法加载");
        }

        BufferedReader reader = null;
        try {
            reader = new BufferedReader(new InputStreamReader(inputStream));
            String line;

            List<String> lines = new ArrayList<>();

            while ((line = reader.readLine()) != null) {
                lines.add(line);
            }

            return loadCustomDnsByLines(lines);

        } finally {
            if (null != reader) {
                try {
                    reader.close();
                } catch (IOException ignored) {
                }
            }
        }

    }

    public static Map<String, List<String>> getCustomDnsMap() {
        return isCustomDnsEnabled() ? copyDomainIpMap(DOMAIN_IP_MAP) : new HashMap<String, List<String>>(0);
    }

    /**
     * 加载自定义的 DNS，要求输入的文件格式是一行表是一个 IP 域名， 可以使用 # 注释
     * <p>
     * IP DOMAIN
     * # 注释
     *
     * @param lines IP 域名列表
     */
    public synchronized static Map<String, List<String>> loadCustomDnsByLines(List<String> lines) {

        Map<String, SortedSet<String>> domainIpMap = buildCustomDnsMapByLines(lines);

        if (null != domainIpMap && !domainIpMap.isEmpty()) {
            LAST_OVERRIDE_CUSTOM_DOMAINS = new HashSet<>();
            LAST_OVERRIDE_CUSTOM_DOMAINS.addAll(DOMAIN_IP_MAP.keySet());

            // 重置
            DOMAIN_IP_MAP = domainIpMap;
            enabledCustomDns();
        }

        return copyDomainIpMap(domainIpMap);
    }

    private static Map<String, List<String>> copyDomainIpMap(Map<String, SortedSet<String>> domainIpMap) {

        Map<String, List<String>> map = new HashMap<>();

        if (null != domainIpMap && !domainIpMap.isEmpty()) {

            for (Map.Entry<String, SortedSet<String>> entry : domainIpMap.entrySet()) {

                String domain = entry.getKey();
                SortedSet<String> ips = entry.getValue();
                if (null != ips && !ips.isEmpty()) {
                    map.put(domain, new ArrayList<>(ips));
                }

            }

        }

        return map;
    }

    private static Map<String, SortedSet<String>> buildCustomDnsMapByLines(List<String> lines) {
        Map<String, SortedSet<String>> domainIpMap = new HashMap<>();
        if (null == lines || lines.isEmpty()) {
            return domainIpMap;
        }

        int lineNumber = 0;
        for (String line : lines) {
            lineNumber++;

            if (StringUtils.isBlank(line)) {
                continue;
            }

            String value = line.trim();
            if (value.startsWith("#")) {
                // 过滤掉注释行
                continue;
            }

            String[] array = value.split("\\s+");
            if (array.length != 2) {
                throw new RuntimeException("自定义HOST，第[" + lineNumber + "]配置格式错误: [" + line + "]");
            }
            String ip = array[0];
            String domain = array[1];

            if (StringUtils.isAnyBlank(ip, domain) || !isValidIp(ip)) {
                throw new RuntimeException("自定义HOST，第[" + lineNumber + "]配置格式错误: [" + line + "]");
            }

            SortedSet<String> set = domainIpMap.get(domain);
            if (set == null) {
                set = new TreeSet<>();
                domainIpMap.put(domain, set);
            }
            set.add(ip);
        }

        return domainIpMap;
    }

    /**
     * 检查IP地址格式是否合法
     *
     * @param ip ip地址
     * @return true 表示地址合法，false 标识不合法
     */
    public static boolean isValidIp(String ip) {
        return IPAddressUtil.isIPv4LiteralAddress(ip) || IPAddressUtil.isIPv6LiteralAddress(ip);
    }

    /**
     * 是否是 IP v4
     *
     * @param ip ip 地址
     * @return true -是， false-否
     */
    public static boolean isIPv4(String ip) {
        return IPAddressUtil.isIPv4LiteralAddress(ip);
    }

    /**
     * 是否是 IP v6
     *
     * @param ip ip 地址
     * @return true -是， false-否
     */
    public static boolean isIPv6(String ip) {
        return IPAddressUtil.isIPv6LiteralAddress(ip);
    }

    /**
     * 获取域名解析到的第一个IP地址
     *
     * @param domain    域名
     * @param defaultIp 默认IP，如果解析错误或者没有对应的IP则使用这个IP
     * @return 返回第一个解析到的IP地址，如果为空表示无法解析
     */
    public static String getFirstIp(String domain, String defaultIp) {
        try {
            InetAddress address = InetAddress.getByName(domain);
            if (null == address) {
                return defaultIp;
            }
            return address.getHostAddress();
        } catch (UnknownHostException e) {
            // 无法解析域名
            return defaultIp;
        }
    }

    /**
     * 获取域名解析到的第一个IP地址
     *
     * @param domain 域名
     * @return 返回第一个解析到的IP地址，如果为空表示无法解析
     */
    public static String getFirstIp(String domain) {
        return getFirstIp(domain, null);
    }

    /**
     * 获取指定域名解析的IP地址列表
     *
     * @param domain 域名
     * @return 返回解析的IP地址，始终返回非null
     */
    public static Set<String> getResolveIps(String domain) {

        Set<String> ips = new HashSet<>();

        try {
            InetAddress[] address = InetAddress.getAllByName(domain);
            if (null == address || address.length < 1) {
                return ips;
            }
            for (InetAddress inetAddress : address) {
                ips.add(inetAddress.getHostAddress());
            }
        } catch (UnknownHostException e) {
            // 无法解析域名
            return ips;
        }
        return ips;
    }
}
